# 开发者：郭同源
# 开发时间：2021/10/30 17:19
import os,cv2

import numpy as np

from dataLoader import MyData
import torch
from my_model import MyModel


if __name__ == '__main__':

    root = 'data/train/train'

    img_name = os.listdir(root)
    for i in img_name:
        img_addr = os.path.join(root,i)
        img = cv2.imread(img_addr)
        info_list = i.split('.')
        label = info_list[1]
        position = []
        sort = info_list[6]

        for j in info_list[2:6]:
            position.append(int(j))

        cv2.rectangle(img,(position[0],position[1]),(position[2],position[3]),color=(0,0,255),thickness=1)
        cv2.putText(img,sort,(position[0],position[1]-3),cv2.FONT_HERSHEY_SIMPLEX,2,color=(0,0,255),thickness=1)

        model = MyModel()
        model.load_state_dict(torch.load('model/model5.pth'))

        new_img = torch.tensor(img,dtype=torch.float32).permute((2,0,1))
        new_img = torch.unsqueeze(new_img,dim=0)/255

        out_label,out_position,out_sort = model(new_img)
        out_sort = torch.argmax(torch.softmax(out_sort, dim=1))


        out_label = torch.sigmoid(out_label)
        out_position = out_position[0]*300
        out_position = [int(i) for i in out_position]





        cv2.rectangle(img,(out_position[0],out_position[1]),(out_position[2],out_position[3]),(0,255,0),1)
        cv2.putText(img,str(out_sort.item()),(out_position[0],out_position[1]-3),
                    cv2.FONT_HERSHEY_SIMPLEX,2,(0,255,0),2)
        cv2.imshow('pic',img)
        cv2.waitKey(300)

        cv2.imwrite('run/output_{}'.format(i),img)







